import re
import numpy as np
import sys
import os

def analyze_step_time(log_file_path):
    """
    从训练日志文件中分析平均单步耗时（剔除异常值）
    """
    # 检查文件是否存在
    if not os.path.isfile(log_file_path):
        print(f"错误：文件 '{log_file_path}' 不存在。")
        return

    pattern = r'^Epoch: \[\d+\]\[\d+/\d+\].*Time ([\d.]+)'

    times = []

    try:
        with open(log_file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                # 只处理以 Epoch 开头的行
                if not line.startswith("Epoch:"):
                    continue

                match = re.match(pattern, line)
                if match:
                    try:
                        t = float(match.group(1))
                        times.append(t)
                    except ValueError:
                        print(f"警告：第 {line_num} 行耗时无法解析: {line}")
    except Exception as e:
        print(f"读取文件时出错: {e}")
        return

    if not times:
        print("未从日志中提取到任何有效的单步耗时数据。")
        return

    times = np.array(times)
    print(f"共提取到 {len(times)} 个 step 的耗时数据。")

    # 原始统计
    print(f"原始耗时范围: {times.min():.4f} ~ {times.max():.4f}")
    print(f"原始平均耗时: {times.mean():.4f} 秒")

    # 使用 IQR 方法剔除异常值
    Q1 = np.percentile(times, 25)
    Q3 = np.percentile(times, 75)
    IQR = Q3 - Q1
    lower_bound = Q1 - 1.5 * IQR
    upper_bound = Q3 + 1.5 * IQR

    filtered = times[(times >= lower_bound) & (times <= upper_bound)]
    outlier_count = len(times) - len(filtered)

    print(f"四分位距 (IQR): {IQR:.4f}")
    print(f"异常值范围: < {lower_bound:.4f} 或 > {upper_bound:.4f}")
    print(f"检测并剔除 {outlier_count} 个异常值")

    if len(filtered) == 0:
        print("警告：所有数据都被识别为异常值，可能是数据分布异常或日志不完整。")
        print(f"建议检查日志，当前数据点: {sorted(times.tolist())}")
    else:
        mean_time = filtered.mean()
        std_time = filtered.std()
        print(f"有效数据范围: {filtered.min():.4f} ~ {filtered.max():.4f}")
        print(f"平均单步耗时（剔除异常后）: {mean_time:.4f} 秒")
        print(f"标准差: {std_time:.4f} 秒")

    # 可选：输出排序后的耗时，便于人工检查
    # print("所有提取的耗时（已排序）:", sorted(times))


if __name__ == "__main__":
    # 如果命令行传入文件路径，使用它；否则提示输入
    if len(sys.argv) < 2:
        print("用法: python analyze_training_time.py <日志文件路径>")
        print("示例: python analyze_training_time.py train.log")
        log_path = input("请输入日志文件路径: ").strip()
    else:
        log_path = sys.argv[1].strip()

    analyze_step_time(log_path)